#pragma once

#include <string>
#include <vector>
#include <map>
#include <memory>
#include <utility>
#include "ipu_model_wrapper.hpp"
#include "ipu_image_preprocessor.hpp"

namespace protonet {

/**
 * @brief 距离度量类型
 */
enum class DistanceMetric {
    EUCLIDEAN,    ///< 欧几里得距离
    COSINE        ///< 余弦相似度
};

/**
 * @brief 置信度计算方法
 */
enum class ConfidenceMethod {
    SIMPLE,                      ///< 简单方法
    NEGATIVE_LOG_LIKELIHOOD      ///< 负对数似然
};

/**
 * @brief 分类结果结构体
 */
struct ClassificationResult {
    std::string predicted_class;      ///< 预测类别
    std::string actual_class;         ///< 实际类别（用于评估）
    float confidence;                 ///< 置信度
    float distance;                   ///< 距离
    std::string image_path;           ///< 图像路径
    bool is_correct;                  ///< 是否预测正确
    DistanceMetric distance_metric;   ///< 使用的距离度量
    ConfidenceMethod confidence_method; ///< 使用的置信度方法
};

/**
 * @brief IPU ProtoNet分类器主类
 */
class IPUProtoNetClassifier {
public:
    /**
     * @brief 构造函数
     * @param model_path IPU模型文件路径
     * @param input_size 输入图像尺寸
     * @param preprocess_config 预处理配置
     */
    IPUProtoNetClassifier(const std::string& model_path, 
                         const std::pair<int, int>& input_size = {224, 224},
                         const IPUImagePreprocessor::PreprocessConfig* preprocess_config = nullptr);
    
    /**
     * @brief 析构函数
     */
    ~IPUProtoNetClassifier() = default;
    
    /**
     * @brief 禁用拷贝构造和赋值
     */
    IPUProtoNetClassifier(const IPUProtoNetClassifier&) = delete;
    IPUProtoNetClassifier& operator=(const IPUProtoNetClassifier&) = delete;
    
    /**
     * @brief 构建类别原型
     * @param dataset_path 数据集路径
     * @return 是否成功
     */
    bool buildPrototypes(const std::string& dataset_path);
    
    /**
     * @brief 对单张图像进行分类
     * @param image_path 图像路径
     * @param distance_metric 距离度量方法
     * @param confidence_method 置信度计算方法
     * @return 分类结果
     */
    ClassificationResult classify(const std::string& image_path,
                                DistanceMetric distance_metric = DistanceMetric::EUCLIDEAN,
                                ConfidenceMethod confidence_method = ConfidenceMethod::SIMPLE);
    
    /**
     * @brief 评估整个数据集
     * @param eval_path 评估数据集路径
     * @param distance_metric 距离度量方法
     * @param confidence_method 置信度计算方法
     * @return 评估结果列表
     */
    std::vector<ClassificationResult> evaluateDataset(const std::string& eval_path,
                                                     DistanceMetric distance_metric = DistanceMetric::EUCLIDEAN,
                                                     ConfidenceMethod confidence_method = ConfidenceMethod::SIMPLE);
    
    /**
     * @brief 保存原型特征到文件
     * @param save_path 保存路径
     * @return 是否成功
     */
    bool savePrototypes(const std::string& save_path) const;
    
    /**
     * @brief 从文件加载原型特征
     * @param load_path 加载路径
     * @return 是否成功
     */
    bool loadPrototypes(const std::string& load_path);
    
    /**
     * @brief 获取类别名称列表
     * @return 类别名称列表
     */
    std::vector<std::string> getClassNames() const { return class_names_; }
    
    /**
     * @brief 获取原型特征数量
     * @return 原型特征数量
     */
    size_t getPrototypeCount() const { return prototypes_.size(); }
    
    /**
     * @brief 检查原型是否已构建
     * @return 是否已构建
     */
    bool hasPrototypes() const { return !prototypes_.empty(); }

private:
    /**
     * @brief 提取图像特征
     * @param image_path 图像路径
     * @return 特征向量
     */
    std::vector<float> extractFeatures(const std::string& image_path);
    
    /**
     * @brief 计算原型特征（多个特征的平均值）
     * @param features_list 特征列表
     * @return 原型特征
     */
    std::vector<float> computePrototype(const std::vector<std::vector<float>>& features_list);
    
    /**
     * @brief 计算两个特征向量间的距离
     * @param features1 第一个特征向量
     * @param features2 第二个特征向量
     * @param metric 距离度量方法
     * @return 距离值
     */
    float computeDistance(const std::vector<float>& features1,
                         const std::vector<float>& features2,
                         DistanceMetric metric);
    
    /**
     * @brief 计算置信度
     * @param distances 距离列表（类别名，距离）
     * @param method 置信度计算方法
     * @return 置信度值
     */
    float computeConfidence(const std::vector<std::pair<std::string, float>>& distances,
                           ConfidenceMethod method);

private:
    std::unique_ptr<IPUModelWrapper> model_wrapper_;          ///< 模型包装器
    std::unique_ptr<IPUImagePreprocessor> image_preprocessor_; ///< 图像预处理器
    std::pair<int, int> input_size_;                          ///< 输入图像尺寸
    
    std::vector<std::string> class_names_;                    ///< 类别名称列表
    std::map<std::string, std::vector<float>> prototypes_;    ///< 原型特征映射
};

} // namespace protonet